#include "render.h"

#include "log.h"

#include <vector>
#include <algorithm>
#include <map>
#include <set>
#include <atlbase.h>

//const int MAX_SURFACES = 10;
const int MAX_SURFACES = 1024 * 10;
static WorldSurface s_surfaces[MAX_SURFACES];
static WorldSurface* s_surfacePtrs[MAX_SURFACES];

const int MAX_INDEXES_PER_BUFFER = 128 * 1024;

int g_surfaceIndex = 0;
bool g_maxSurfacesExceeded = false;

CComPtr<ID3D11Device> g_device = NULL;
CComPtr<ID3D11Buffer> g_surfIndexBuffer = NULL;
CComPtr<ID3D11PixelShader> g_ps1Stage = NULL;
CComPtr<ID3D11PixelShader> g_ps2Stage = NULL;
CComPtr<ID3D11PixelShader> g_ps3Stage = NULL;
CComPtr<ID3D11PixelShader> g_ps4Stage = NULL;
std::vector<ID3D11ShaderResourceView*> g_lightmaps;
std::map<std::string, RenderShader> g_renderShaders;


namespace {
// TODO only 1 copy of this
HRESULT CompileShaderFromFile(WCHAR* szFileName, D3D_SHADER_MACRO* macros, LPCSTR szEntryPoint, LPCSTR szShaderModel, ID3DBlob** ppBlobOut)
{
    HRESULT hr = S_OK;

    DWORD dwShaderFlags = D3DCOMPILE_ENABLE_STRICTNESS;
#if defined(DEBUG) || defined(_DEBUG)
    // Set the D3DCOMPILE_DEBUG flag to embed debug information in the shaders.
    // Setting this flag improves the shader debugging experience, but still allows 
    // the shaders to be optimized and to run exactly the way they will run in 
    // the release configuration of this program.
    dwShaderFlags |= D3DCOMPILE_DEBUG;
#endif

    ID3DBlob* pErrorBlob;
    hr = D3DX11CompileFromFile(szFileName, macros, NULL, szEntryPoint, szShaderModel, 
        dwShaderFlags, 0, NULL, ppBlobOut, &pErrorBlob, NULL);
    if(FAILED(hr))
    {
        if(pErrorBlob != NULL)
		{
			MessageBoxA(NULL, (char*)pErrorBlob->GetBufferPointer(), "Compile Error", MB_OK | MB_SYSTEMMODAL); 
		}
       //     OutputDebugStringA((char*)pErrorBlob->GetBufferPointer());
        if(pErrorBlob) pErrorBlob->Release();
        return hr;
    }
    if(pErrorBlob) pErrorBlob->Release();

    return S_OK;
}

HRESULT LoadPixelShader(const std::string& name, ID3D11PixelShader** shader)
{
	if (shader == NULL)
	{
		return E_FAIL;
	}
		
	*shader = NULL;

	ID3DBlob* pixelShaderBlob = NULL;
	HRESULT hr = CompileShaderFromFile(L"Tutorial02.fx", NULL, name.c_str(), "ps_4_0", &pixelShaderBlob);
	if (FAILED(hr))
	{
		return hr;
	}

	hr = g_device->CreatePixelShader(pixelShaderBlob->GetBufferPointer(), pixelShaderBlob->GetBufferSize(),
		NULL, shader);
	pixelShaderBlob->Release();
	if (FAILED(hr))
	{
		return hr;
	}

	return S_OK;
}

}

HRESULT Render_Init(ID3D11Device* device, int maxIndexes)
{
	g_device = device;
	assert(g_device != NULL);

	Log_Printf("maxIndexes=%d\n", maxIndexes);

	D3D11_BUFFER_DESC bd;
	ZeroMemory(&bd, sizeof(bd));
	bd.Usage = D3D11_USAGE_DYNAMIC;
	bd.BindFlags = D3D11_BIND_INDEX_BUFFER;
	bd.ByteWidth = MAX_INDEXES_PER_BUFFER * sizeof(int);
	bd.CPUAccessFlags = D3D11_CPU_ACCESS_WRITE;
	HRESULT hr = device->CreateBuffer(&bd, NULL, &g_surfIndexBuffer);
	if (FAILED(hr))
	{
		__debugbreak();
	}

	if (FAILED(LoadPixelShader("shader_1", &g_ps1Stage)) ||
		FAILED(LoadPixelShader("shader_2", &g_ps2Stage)) ||
		FAILED(LoadPixelShader("shader_3", &g_ps3Stage)) ||
		FAILED(LoadPixelShader("shader_4", &g_ps4Stage)))
	{
		return E_FAIL;
	}

	return S_OK;
}

void Render_Shutdown()
{
	for (auto iter = g_renderShaders.begin(); iter != g_renderShaders.end(); ++iter)
	{
		for (auto stageIter = iter->second.stages.begin(); stageIter != iter->second.stages.end(); ++stageIter)
		{
			SAFE_RELEASE((*stageIter).texture);
		}
	}

	/*
	SAFE_RELEASE(g_ps1Stage);
	SAFE_RELEASE(g_ps2Stage);
	SAFE_RELEASE(g_ps3Stage);
	SAFE_RELEASE(g_ps4Stage);

	SAFE_RELEASE(g_surfIndexBuffer);
	SAFE_RELEASE(g_device);
	*/
}

bool FileExists(const std::string& path)
{
	FILE* f = fopen(path.c_str(), "r");
	if (f == NULL)
	{
		return false;
	}

	fclose(f);
	return true;
}

size_t FileSize(FILE* f)
{
	const size_t pos = ftell(f);
	fseek(f, 0, SEEK_END);
	const size_t size = ftell(f);
	fseek(f, pos, SEEK_SET);
	return size;
}

template <typename T>
struct PtrHolder
{
	PtrHolder(T* p)
		: ptr_(p)
	{
	}

	virtual ~PtrHolder()
	{
		delete ptr_;
	}

	operator T*()
	{
		return ptr_;
	}

	T*		ptr_;
};

template <typename T>
struct ArrayHolder
{
	ArrayHolder(T* p)
		: ptr_(p)
	{
	}

	virtual ~ArrayHolder()
	{
		delete[] ptr_;
	}

	operator T*()
	{
		return ptr_;
	}

	T*		ptr_;
};

bool Render_CreateLightmapImage(const unsigned char* data, const size_t size)
{
	ID3D11ShaderResourceView* resource = NULL;

	/*
	D3DX11_IMAGE_LOAD_INFO info;
	ZeroMemory(&info, sizeof(info));
	info.Width = 128;
	info.Height = 128;
	info.Depth = 24;
	info.Format = dxgi_form
		*/

	HRESULT hr = D3DX11CreateShaderResourceViewFromMemory(g_device, data, size, NULL, NULL, &resource, NULL);
	if (FAILED(hr))
	{
		return false;
	}

	g_lightmaps.push_back(resource);
	return true;
}

ID3D11ShaderResourceView* CreateShaderResourceViewFromTGA(const std::string& path)
{
	FILE* f = fopen(path.c_str(), "rb");
	if (f == NULL)
	{
		return NULL;
	}

	const size_t fileSize = FileSize(f);
	ArrayHolder<BYTE> buffer(new BYTE[fileSize]);
	const size_t read = fread(buffer, 1, fileSize, f);
	fclose(f);
	if (read != fileSize)
	{
		Log_Printf("Error: couldn't read file '%s'\n", path.c_str());
		return NULL;
	}

	struct TGAHeader
	{
		BYTE		idLength;
		BYTE		colourMapType;
		BYTE		imageType;
		BYTE		colourMapSpec[5];
		short		xOrigin;
		short		yOrigin;
		short		width;
		short		height;
		BYTE		depth;
		BYTE		descriptor;
	};
	TGAHeader* hdr = reinterpret_cast<TGAHeader*>((BYTE*)buffer);

	const int TYPE_UNCOMPRESSED_COLOUR = 2;
	if (hdr->imageType != TYPE_UNCOMPRESSED_COLOUR)
	{
		Log_Printf("Error: unsupported TGA type '%d' for file '%s'\n", hdr->imageType, path.c_str());
		return NULL;
	}

	D3D11_TEXTURE2D_DESC desc;
	desc.Width = hdr->width;
	desc.Height = hdr->height;
	desc.ArraySize = 1;
	desc.SampleDesc.Count = 1;
	desc.SampleDesc.Quality = 0;
	desc.Format = DXGI_FORMAT_R8G8B8A8_UNORM;
	desc.Usage = D3D11_USAGE_DEFAULT;
	desc.BindFlags = D3D11_BIND_RENDER_TARGET | D3D11_BIND_SHADER_RESOURCE;
	desc.CPUAccessFlags = 0;
	desc.MiscFlags = D3D11_RESOURCE_MISC_GENERATE_MIPS;
	desc.MipLevels = 0;

	D3D11_SUBRESOURCE_DATA initData;
	initData.pSysMem = buffer + sizeof(TGAHeader);
	initData.SysMemPitch = (hdr->depth/8) * hdr->width;

	ID3D11Texture2D* d3dTexture = NULL;
	HRESULT hr = g_device->CreateTexture2D(&desc, NULL, &d3dTexture);
	if (FAILED(hr))
	{
		Log_Printf("Error: could not create Texture2D (0x%08x)\n", hr);
		return NULL;
	}

	ID3D11DeviceContext* context = NULL;
	g_device->GetImmediateContext(&context);
	if (context == NULL)
	{
		Log_Printf("Error: could not get immediate context", hr);
		return NULL;
	}

	context->UpdateSubresource(d3dTexture, D3D11CalcSubresource(0, 0, 0), NULL,
		buffer, (hdr->depth/8) * hdr->width, 0);

	D3D11_TEXTURE2D_DESC tDesc;
	d3dTexture->GetDesc(&tDesc);

	D3D11_SHADER_RESOURCE_VIEW_DESC rvDesc;
	rvDesc.Format = tDesc.Format;
	rvDesc.ViewDimension = D3D11_SRV_DIMENSION_TEXTURE2D;
	rvDesc.Texture2D.MostDetailedMip = 0;
	rvDesc.Texture2D.MipLevels = tDesc.MipLevels;

	ID3D11ShaderResourceView* resourceView = NULL;
	if (FAILED(g_device->CreateShaderResourceView(d3dTexture, &rvDesc, &resourceView)))
	{
		Log_Printf("Error: could not create shader resource view for texture '%s'\n", path.c_str());
		return NULL;
	}

	return resourceView;
}

bool Render_CreateShader(const Shader& shader)
{
	RenderShader renderShader;
	for (auto iter = shader.stages.begin(); iter != shader.stages.end(); ++iter)
	{
		const ShaderStage& stage = *iter;
		RenderShaderStage renderShaderStage;
		renderShader.name = shader.name;

		if (stage.texture == "$lightmap")
		{
			// TODO a better way of communicating 'lightmap'
			renderShaderStage.texture = NULL;
		}
		else 
		{
			ID3D11ShaderResourceView* resourceView = NULL;
			const std::string& path = stage.texture;
			if (path.length() == 0)
			{
				return false;
			}

			const std::string extension = path.substr(path.length() - 3);
			if (extension == "tga" && FileExists(path))
			{
				resourceView = CreateShaderResourceViewFromTGA(path);
			}
			else
			{
				std::string jpgPath = path.substr(0, path.length() - 3) + std::string("jpg");
				if (FileExists(jpgPath))
				{
					HRESULT hr = D3DX11CreateShaderResourceViewFromFileA(g_device,
						jpgPath.c_str(), NULL, NULL, &resourceView, NULL);
					if (FAILED(hr))
					{
						return false;
					}
				}
			}

			if (resourceView != NULL)
			{
				renderShaderStage.texture = resourceView;
			}
			else
			{
				// TODO set default texture
			}
		}

		renderShaderStage.srcParam = stage.srcParam;
		renderShaderStage.dstParam = stage.dstParam;
		if (renderShaderStage.texture == (ID3D11ShaderResourceView*)0xcccccccc)
			__debugbreak();

		renderShader.stages.push_back(renderShaderStage);
	}

	if (renderShader.stages.size() > 0)
	{
	//	Log_Printf("Loaded shader '%s'\n", shader.name.c_str());
		g_renderShaders[shader.name] = renderShader;
	}

	return true;
}

bool Render_CreateShaderFromTexturePath(const std::string& path)
{
	if (g_renderShaders.find(path) != g_renderShaders.end())
	{
		// already a shader loaded with this name
		return true;
	}

	// try loading just the texture
	const char* extensions[] = { "tga", "jpg", "jpeg", NULL };
	for (int e = 0; extensions[e] != NULL; e++)
	{
		char buffer[256];
		sprintf(buffer, "data/baseq3/%s.%s", path.c_str(), extensions[e]);

		ID3D11ShaderResourceView* resourceView = NULL;
		if (strcmp(extensions[e], "tga") == 0)
		{
			resourceView = CreateShaderResourceViewFromTGA(buffer);
			if (resourceView == NULL)
			{
				continue;
			}
		}
		else
		{
			HRESULT hr = D3DX11CreateShaderResourceViewFromFileA(g_device, buffer,
				NULL, NULL, &resourceView, NULL);
			if (FAILED(hr) || resourceView == NULL)
			{
				continue;
			}
		}

		RenderShader shader;
		RenderShaderStage stage;
		stage.dstParam = Blend_One;
		stage.srcParam = Blend_One;
		stage.texture = resourceView;
		shader.stages.push_back(stage);

		// add a lightmap stage too
		stage.texture = NULL;
		shader.stages.push_back(stage);
		shader.name = path;
	
		g_renderShaders[path] = shader;
		return true;
	}

	return false;
}

void Render_BeginFrame()
{
	g_surfaceIndex = 0;
	g_maxSurfacesExceeded = false;
}

WorldSurface* Render_AllocSurface()
{
	if (g_surfaceIndex >= MAX_SURFACES)
	{
		if (!g_maxSurfacesExceeded)
		{
			Log_Printf("Max surfaces (%d) exceeded!\n", MAX_SURFACES);
			g_maxSurfacesExceeded = true;
		}
	
		return NULL;
	}

	s_surfacePtrs[g_surfaceIndex] = &s_surfaces[g_surfaceIndex];
	return s_surfacePtrs[g_surfaceIndex++];
}

WorldSurface* Render_GetFirstSurface()
{
	return s_surfacePtrs[0];
}

int	Render_GetNumSurfaces()
{
	return g_surfaceIndex;
}

struct surface_compare
{
	bool operator()(const WorldSurface* lhs, const WorldSurface* rhs)
	{
		// the presence or not of a lightmap determines the pixel shader, sort by that first
		if (lhs->lightmap < rhs->lightmap)
		{
			return true;
		}
		else if (lhs->lightmap == rhs->lightmap)
		{
			// finally group by texture
			return lhs->texture < rhs->texture;
		}

		return false;
	}
};
#if 0
int indexes[MAX_INDEXES_PER_BUFFER];
int BuildIndexBufferData(int surfType)
{
	WorldSurface* worldSurfaces = Render_GetFirstSurface();
	const int numWorldSurfaces = Render_GetNumSurfaces();
//	Log_Printf("surfaces=%d\n", numWorldSurfaces);

	int count = 0;
	int* p = &indexes[0];
	for (int i = 0; i < numWorldSurfaces; i++)
	{
		WorldSurface* s = worldSurfaces + i;

		if ((s->type & surfType) == 0)
			continue;

		memcpy(p, &s->indexes[0], sizeof(int) * s->numIndexes);
		p += s->numIndexes;
		count += s->numIndexes;
	}

	ID3D11DeviceContext* context = NULL;
	g_device->GetImmediateContext(&context);
	if (context != NULL)
	{
		D3D11_MAPPED_SUBRESOURCE resource;
		HRESULT hr = context->Map(g_surfIndexBuffer, 0, D3D11_MAP_WRITE_DISCARD, 0, &resource);
		memcpy(resource.pData, indexes, sizeof(int) * count);
		context->Unmap(g_surfIndexBuffer, NULL);
	}

	context->IASetIndexBuffer(g_surfIndexBuffer, DXGI_FORMAT_R32_UINT, 0);
	return count;
}
#endif

int BuildIndexBufferForSurfaces(const std::vector<WorldSurface*>& surfaces, ShaderSurfaces& renderSurfaces, int offset)
{
	const int MAX_INDEXES_PER_BUFFER = 1024 * 50;
	int indexes[MAX_INDEXES_PER_BUFFER];

	int count = 0;
	for (auto iter = surfaces.begin(); iter != surfaces.end(); ++iter)
	{
		WorldSurface* s = *iter;
		if (count + s->numIndexes < MAX_INDEXES_PER_BUFFER)
		{
			memcpy(&indexes[count], &s->indexes[0], sizeof(int) * s->numIndexes);
			count += s->numIndexes;
		}
	}

	ID3D11DeviceContext* context = NULL;
	g_device->GetImmediateContext(&context);
	if (context == NULL)
	{
		return 0;
	}

	D3D11_MAPPED_SUBRESOURCE resource;
	D3D11_MAP mapType = D3D11_MAP_WRITE_NO_OVERWRITE;
	if (offset == 0)
	{
		mapType = D3D11_MAP_WRITE_DISCARD;
	}

	HRESULT hr = context->Map(g_surfIndexBuffer, 0, mapType, 0, &resource); 
	int* dataPtr = (int*)resource.pData + offset;
	memcpy(dataPtr, indexes, sizeof(int) * count);
	context->Unmap(g_surfIndexBuffer, 0);

	renderSurfaces.firstIndex = offset;
	renderSurfaces.numIndexes = count;
	return count;
}

void Render_DrawFrame()
{
//	Log_Printf("====== Begin frame ======\n");

	ID3D11DeviceContext* ctx = NULL;
	g_device->GetImmediateContext(&ctx);
	assert(ctx != NULL);

	std::vector<WorldSurface*> drawableSurfaces;
	std::set<RenderShader*> shaders;
	WorldSurface* surfaces = Render_GetFirstSurface();
	for (int i = 0; i < Render_GetNumSurfaces(); i++)
	{
		WorldSurface* s = surfaces + i;
		if (g_renderShaders.find(s->texture) != g_renderShaders.end())
		{
			drawableSurfaces.push_back(s);
			shaders.insert(&g_renderShaders[s->texture]);
		}
	}

	std::map<int, std::vector<WorldSurface*>> surfaceBuckets;
	for (auto iter = drawableSurfaces.begin(); iter != drawableSurfaces.end(); ++iter)
	{
		WorldSurface* s = *iter;
		RenderShader* shader = &g_renderShaders[s->texture];

		const int numStages = shader->stages.size();
//		Log_Printf("t=%s num=%d\n", s->texture.c_str(), numStages);
		if (surfaceBuckets.find(numStages) == surfaceBuckets.end())
		{
			std::vector<WorldSurface*> shaderSurfs;
			shaderSurfs.push_back(s);
			surfaceBuckets.insert(std::pair<int, std::vector<WorldSurface*>>(numStages, shaderSurfs));
		}
		else
		{
			surfaceBuckets[numStages].push_back(s);
		}
	}

	Scene scene;
	int indexCount = 0;
	for (auto iter = surfaceBuckets.begin(); iter != surfaceBuckets.end(); ++iter)
	{
		std::vector<WorldSurface*>& surfs = iter->second;
		std::sort(surfs.begin(), surfs.end(), surface_compare());
	
		PixelShaderGroup shaderGroup;
		ShaderSurfaces renderSurfs;
		std::vector<WorldSurface*> shaderSurfs;
		for (auto surfIter = surfs.begin(); surfIter != surfs.end(); ++surfIter)
		{
			WorldSurface* s = *surfIter;
			if (renderSurfs.lightmap != s->lightmap ||
				renderSurfs.shader != s->texture)
			{
				// a new group of surfaces to be rendered together
				if (renderSurfs.shader != "default")
				{
					indexCount += BuildIndexBufferForSurfaces(shaderSurfs, renderSurfs, indexCount);
					shaderGroup.surfaces.push_back(renderSurfs);
				}

				renderSurfs = ShaderSurfaces();
				renderSurfs.shader = s->texture;
				renderSurfs.lightmap = s->lightmap;
				shaderSurfs.clear();
			}

			shaderSurfs.push_back(s);
		}

		if (shaderSurfs.size() > 0)
		{
			indexCount += BuildIndexBufferForSurfaces(shaderSurfs, renderSurfs, indexCount);
			shaderGroup.surfaces.push_back(renderSurfs);
		}

		scene.psGroups.push_back(shaderGroup);
	}

	ID3D11DeviceContext* context = NULL;
	g_device->GetImmediateContext(&context);
	if (context == NULL)
	{
		return;
	}

	context->IASetIndexBuffer(g_surfIndexBuffer, DXGI_FORMAT_R32_UINT, 0);
	for (auto iter = scene.psGroups.begin(); iter != scene.psGroups.end(); ++iter)
	{
		PixelShaderGroup& group = *iter;

		ShaderSurfaces& firstSurfs = group.surfaces.front();
		assert(g_renderShaders.find(firstSurfs.shader) != g_renderShaders.end());
		RenderShader& firstShader = g_renderShaders[firstSurfs.shader];
		const int numPixelShaderStages = firstShader.stages.size();
		switch (numPixelShaderStages)
		{
		case 1:
			context->PSSetShader(g_ps1Stage, NULL, 0);
			break;
		case 2:
			context->PSSetShader(g_ps2Stage, NULL, 0);
			break;
		case 3:
			context->PSSetShader(g_ps3Stage, NULL, 0);
			break;
		case 4:
			context->PSSetShader(g_ps4Stage, NULL, 0);
			break;
		}

		for (auto surfIter = group.surfaces.begin(); surfIter != group.surfaces.end(); ++surfIter)
		{
			ShaderSurfaces& ss = *surfIter;
			RenderShader& rs = g_renderShaders[ss.shader];
			const size_t numStages = rs.stages.size();
			assert(numStages == numPixelShaderStages);
			for (size_t i = 0; i < numStages; i++)
			{
				if (rs.stages[i].texture == NULL)	// i.e. lightmap
				{
					if (ss.lightmap != -1)
					{
						// TODO this is a shit way of always getting the lightmap into texture0
						assert(ss.lightmap < g_lightmaps.size());
						context->PSSetShaderResources(0, 1, &g_lightmaps[ss.lightmap]);
					}

					break;
				}
			}

			for (size_t i = 0; i < numStages; i++)
			{
				if (rs.stages[i].texture != NULL)
				{
					context->PSSetShaderResources(i + 1, 1, &rs.stages[i].texture);
				}
			}

			context->DrawIndexed(ss.numIndexes, ss.firstIndex, 0);
		}
	}

//	Log_Printf("====== End frame ======\n");
}
